import json
import torch.utils.data as data
import os
from pycocotools.coco import COCO
import cv2 as cv
import numpy as np
from tqdm import tqdm
from data_enhancement.resize_convert import ResizeConvert
from utils.util import draw_umich_gaussian
from matplotlib import pyplot as plt

        

class CocoDataset(data.Dataset):
    
    def __init__(self):
        self.image_dir = r'I:\Datasets\val2017'
        self.config_json = r'I:\Datasets\annotations\person_keypoints_val2017.json'
        # 创建coco 数据集
        self.coco = COCO(self.config_json)
        # 获取所有分类是人的id
        self.persons = self.coco.getCatIds(catNms=['person'])
        #获取所有的图片id
        self.imgIds = self.coco.getImgIds(catIds=self.persons)
 
            
        
        #骨骼宽度百分比
        self.bone_width = 0.0001
        
        #人体关键点的类别和在数组中的位置
        self.keypoint_names = [
            [0, 'nose', '鼻子'],
            [1, 'left_ear', '左耳'],
            [2, 'right_ear', '右耳'],
            [3, 'left_shoulder', '左肩'],
            [4, 'right_shoulder', '右肩'],
            [5, 'neck'  , '脖子'],
            [6, 'left_elbow', '左手肘'],
            [7, 'right_elbow', '右手肘'],
            [8, 'left_wrist', '左手腕'],
            [9, 'right_wrist', '右手腕'],
            [10, 'left_hip', '左臀部'],
            [11, 'right_hip', '右臀部'],
            [12, 'left_knee', '左膝盖'],
            [13, 'right_knee', '右膝盖'],
            [14, 'left_ankle', '左脚踝'],
            [15, 'right_ankle', '右脚踝'],
        ]
        # 骨骼连线
        self.bone_links = {
            0 : [
                1, 2, 5
            ],
            5: [
                3, 4, 10, 11
            ],
            3 : [
                6      
            ],
            6 : [
                8
            ],
            4 : [
                7
            ],
            7 : [
                9
            ],
            10 : [
                12
            ],
            12 : [
                14
            ],
            11 : [
                13
            ],
            13 : [
                15
            ]
        }
        # 点数量
        self.map_len = len(self.keypoint_names)
        
        self.target_size = (512, 512)
        self.resize = ResizeConvert(self.target_size)
        
        self.target_full_area = self.target_size[0] * self.target_size[1]
        # 筛选数据
        self.__load()
    
    def __load(self):
        # 所有的数据集合
        self.datas = []
        #遍历所有的图片
        for img_id in self.imgIds:
            # 读取图片信息
            img_info = self.coco.loadImgs([img_id])[0]
            #图片宽高
            org_img_width = img_info['width']
            org_img_height = img_info['height']
            # 获取图片路径
            image_path = img_info['file_name']
            #拼接图片路径
            image_path = os.path.join(self.image_dir, image_path)
            #获取改图片所有的annid
            annids = self.coco.getAnnIds(imgIds=[img_id] , catIds=self.persons, iscrowd=None)
            # 读取所有的分类信息
            anns = self.coco.loadAnns(annids)
            item = {'image_path': image_path, 'width': org_img_width, 'height': org_img_height}
            points = []
            for ann in anns:
                if ann['num_keypoints'] < 5:
                    continue
                points.append(self.__get_points(ann))
            
            if len(points) == 0:
                continue
        
            item['points'] = points
            # 数据增加
            self.datas.append(item)
        
              
                
    def __len__(self):
        return len(self.anns)
    
    #获取数组列表
    def __get_points(self, ann):
        keypoints = ann['keypoints']
        # print(keypoints)
        keypoints = np.array(keypoints).reshape(-1, 3)
        points = []
        points.append(keypoints[0])
        points.append(keypoints[3])
        points.append(keypoints[4])
        points.append(keypoints[5])
        points.append(keypoints[6])
        points.append([0, 0, 0] if keypoints[6][2] == 0 or keypoints[5][2] == 0 else [(keypoints[5][0] + keypoints[6][0]) / 2, (keypoints[5][1] + keypoints[6][1]) / 2, keypoints[5][2]])
        points.append(keypoints[7])
        points.append(keypoints[8])
        points.append(keypoints[9])
        points.append(keypoints[10])
        points.append(keypoints[11])
        points.append(keypoints[12])
        points.append(keypoints[13])
        points.append(keypoints[14])
        points.append(keypoints[15])
        points.append(keypoints[16])
        return points
    
    # 绘制骨骼
    def __draw_bones(self, area, points, bones, key, img):
        w = min(max(area * self.bone_width, 3), 10)
        
        for v in bones:
            if points[v][2] == 0:
                continue
            vec = points[v] - points[key]
            rad = np.arctan(vec[1] / vec[0])
            rad = rad / (2 * np.pi) * 255.0
            rad = int(rad)
            # print('points1 ：' + str(points[v]), ', points2 ：' + str(points[key]), ', rad：' + str(rad), ', w：' + str(w))
            cv.line(img, (int(points[key][0]), int(points[key][1])), (int(points[v][0]), int(points[v][1])), 255, int(w))
            
    
    
    def __draw_ann(self, ann, img, keypoint_maps):
        keypoints = ann['keypoints']
        num_keypoints = ann['num_keypoints']
        
        if num_keypoints != keypoints / 3:
            return None, None
        
        
        
            
    
    #获取其中一个
    def __getitem__(self, index):
        #获取图片信息
        img_id = self.imgIds[index]
        # 读取图片信息
        img_info = self.coco.loadImgs([img_id])[0]
        #图片宽高
        org_img_width = img_info['width']
        org_img_height = img_info['height']
        # 获取图片路径
        image_path = img_info['file_name']
        #拼接图片路径
        image_path = os.path.join(self.image_dir, image_path)
        # 读取突破数据
        img = cv.imdecode(np.fromfile(image_path, dtype=np.uint8), -1)
        #获取改图片所有的annid
        annids = self.coco.getAnnIds(imgIds=[img_id] , catIds=self.persons, iscrowd=None)
        # 读取所有的分类信息
        anns = self.coco.loadAnns(annids)
        
        #创建关键图片
        keypoint_maps = np.zeros((self.map_len, self.target_size[0], self.target_size[1]), dtype=np.float32)
        # 向量图
        vec_maps = np.zeros((self.map_len, self.target_size[0], self.target_size[1]), dtype=np.uint8)
        points = []
        areas = []
        area_rate = (self.target_size[0] * self.target_size[1]) / (org_img_width * org_img_height)
        for ann in anns:
            points.append(self.__get_points(ann))
            areas.append(ann['area'])
        
        points = np.array(points).astype(np.float32).reshape(-1, 3)
        areas = np.array(areas)
        areas *= area_rate
        
        print(points)
        #转换点
        img, _, points = self.resize.convert(img, None, points)
        img = cv.cvtColor(img, cv.COLOR_BGR2RGB)
        #转换回来
        points = points.reshape(-1, self.map_len, 3)
        #遍历
        for j in range(len(points)):
            p = points[j]
            for index in range(self.map_len):
                if p[index][2] == 0:
                    continue
                radus = int(areas[j] / self.target_full_area * areas[j] * 0.001)
                # 绘制热图
                center = (int(p[index][0]), int(p[index][1]))
                # print('center : ', center, ', radus : ', radus, ', index : ', radus)
                draw_umich_gaussian(keypoint_maps[index], center, 20)
                
                if index in self.bone_links.keys():
                    # print('draw bone!!')
                    self.__draw_bones(areas[j], p, self.bone_links[index], index, vec_maps[index])
            
                cv.circle(img, center, 2, (0, 0, 255), 2)
        
        vec_maps = vec_maps.astype(np.float32) / 255.0
        
        fig = plt.figure()
        for index in range(len(keypoint_maps)):
            plt.subplot(1, 3, 1)
            plt.title('keypoint_map : ' + self.keypoint_names[index][1])
            plt.imshow(keypoint_maps[index], cmap='gray')
            
            plt.subplot(1, 3, 2)
            plt.title('vec_map')
            plt.imshow(vec_maps[index], cmap='gray')
            
            plt.subplot(1, 3, 3)
            plt.title('map')
            plt.imshow(img)
            
            plt.get_current_fig_manager().window.state('zoomed')
            plt.show()
                
        return img, keypoint_maps, vec_maps
        


